import numpy as np
 
# 函数准备
# 激活函数
def gelu(x):
   return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x**3)))

# softmax
def softmax(x):
    exp_x = np.exp(x - np.max(x, axis=-1, keepdims = True))
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

# 层标准化
def layer_norm(x, g, b, eps: float = 1e-5):
    mean = np.mean(x, axis=-1, keepdims=True)
    variance = np.var(x, axis=-1, keepdims=True)
    # 均值为0,方差为1
    x = (x - mean) / np.sqrt(variance + eps)
    return g * x + b

# 线性
def linear(x, w, b):
    return x @ w + b

# 定位的前馈神经网络(2层的简单多层感知, 升维 -> 降维)
def ffn(x, c_fc, c_proj):
    # 升维 -> 激活函数
    a = gelu(linear(x, **c_fc))
    # 降维
    x = linear(a, **c_proj)
    return x

# 多头因果自我关注, begin
# 注意力
def attention(q, k, v, mask):
    return softmax(q @ k.T / np.sqrt(q.shape[-1]) + mask) @ v # 需要在softmax前修改

# 自我 + 因果 + 多头
def mha(x, c_attn, c_proj, n_head):
    x = linear(x, **c_attn)
    qkv = np.split(x, 3, axis=-1)
    qkv_heads = list(map(lambda x: np.split(x, n_head, axis=-1), qkv))
    causal_mask = (1 - np.tri(x.shape[0], dtype=x.dtype)) * -1e10
    # x = attention(q,k,v, causal_mask)
    out_heads = [attention(q,k,v, causal_mask) for q,k,v in zip(*qkv_heads)]
    x = np.hstack(out_heads)
    x = linear(x, **c_proj)
    return x
    
# 多头因果自我关注, end

# 解码器块
def transformer_block(x, mlp, attn, ln_1, ln_2, n_head):
     # 多头因果自我关注
     x = x + mha(layer_norm(x, **ln_1), **attn, n_head=n_head)
     # 定位的前馈神经网络
     x = x + ffn(layer_norm(x, **ln_2), **mlp)
     return x
 
def gpt2(inputs, wte, wpe, blocks, ln_f, n_head):
    # transformer架构
    # 词嵌入矩阵 + 位置嵌入矩阵
    x = wte[inputs] + wpe[range(len(inputs))]
    # 解码器
    for block in blocks:
        x = transformer_block(x, **block, n_head=n_head)
    # 投影
    x = layer_norm(x, **ln_f) # gpt2特有
    return x @ wte.T
 
def generate(inputs, params, n_head, n_tokens_to_generate):
    from tqdm import tqdm
 
 
    for _ in tqdm(range(n_tokens_to_generate), "generating"):  # auto-regressive decode loop
        logits = gpt2(inputs, **params, n_head=n_head)  # model forward pass
        next_id = np.argmax(logits[-1])  # greedy sampling
        inputs.append(int(next_id))  # append prediction to input
 
 
    return inputs[len(inputs) - n_tokens_to_generate :] # only return generated ids
 
 
 
 
def main(prompt: str, n_tokens_to_generate: int = 40, model_size: str = "124M", models_dir: str = "models"):
    from utils import load_encoder_hparams_and_params
 
 
    # load encoder, hparams, and params from the released open-ai gpt-2 files
    encoder, hparams, params = load_encoder_hparams_and_params(model_size, models_dir)
 
 
    # encode the input string using the BPE tokenizer
    input_ids = encoder.encode(prompt)
 
 
    # make sure we are not surpassing the max sequence length of our model
    assert len(input_ids) + n_tokens_to_generate < hparams["n_ctx"]
 
 
    # generate output ids
    output_ids = generate(input_ids, params, hparams["n_head"], n_tokens_to_generate)
 
 
    # decode the ids back into a string
    output_text = encoder.decode(output_ids)
 
 
    return output_text
 
 
 
 
if __name__ == "__main__":
    import fire
 
    fire.Fire(main)